#!/usr/bin/env python3
"""
生成DanceTrack数据集的序列映射文件(seqmap.txt)
根据TrackEval的格式要求，每行仅包含：
<sequence name>
"""

import os
import sys
from configparser import ConfigParser

def generate_seqmap(data_root, split, output_file):
    """
    生成指定split的seqmap文件
    
    Args:
        data_root: 数据集根目录，例如 "./datasets"
        split: 数据集分割，例如 "train", "val", "test"
        output_file: 输出文件路径
    """
    dancetrack_dir = os.path.join(data_root, "DanceTrack", split)
    
    if not os.path.exists(dancetrack_dir):
        print(f"错误：目录不存在 {dancetrack_dir}")
        return False
    
    sequences = []
    sequence_dirs = [d for d in os.listdir(dancetrack_dir) 
                    if os.path.isdir(os.path.join(dancetrack_dir, d))]
    
    for seq_name in sorted(sequence_dirs):
        seq_dir = os.path.join(dancetrack_dir, seq_name)
        seqinfo_file = os.path.join(seq_dir, "seqinfo.ini")
        
        if not os.path.exists(seqinfo_file):
            print(f"警告：未找到 {seqinfo_file}，跳过序列 {seq_name}")
            continue
        
        try:
            # 读取seqinfo.ini文件
            config = ConfigParser()
            config.read(seqinfo_file)
            
            seq_length = int(config["Sequence"]["seqLength"])
            im_height = int(config["Sequence"]["imHeight"])
            im_width = int(config["Sequence"]["imWidth"])
            
            sequences.append((seq_name, seq_length, im_height, im_width))
            print(f"处理序列: {seq_name} - {seq_length}帧, {im_width}x{im_height}")
            
        except Exception as e:
            print(f"错误：处理序列 {seq_name} 时出错: {e}")
            continue
    
    # 写入seqmap文件（每行仅序列名，符合TrackEval的MotChallenge2DBox要求）
    try:
        with open(output_file, 'w') as f:
            for seq_name, _, _, _ in sequences:
                f.write(f"{seq_name}\n")
        
        print(f"成功生成 {output_file}，包含 {len(sequences)} 个序列")
        return True
        
    except Exception as e:
        print(f"错误：写入文件 {output_file} 时出错: {e}")
        return False

def main():
    # 使用固定的数据集路径
    data_root = "/home/liuyonghui/datasets/"
    
    # 为每个split生成seqmap文件
    splits = ["train", "val", "test"]
    
    for split in splits:
        output_file = os.path.join(data_root, "DanceTrack", f"{split}_seqmap.txt")
        print(f"\n正在生成 {split} 的seqmap文件...")
        
        success = generate_seqmap(data_root, split, output_file)
        if not success:
            print(f"生成 {split}_seqmap.txt 失败")
        else:
            print(f"✓ {split}_seqmap.txt 生成成功")

if __name__ == "__main__":
    main()
